require("Rniftilib")

source("base.R")

setOldClass("nifti")

setClass(
    Class = "nifti_array", 
    representation = representation(
        header = "list"
    ),
    contains = "array"
)

# 2D matrix to represent 4D array
setClass(
    Class = "nifti_matrix",
    representation = representation(
        header = "list",
        mask_used = "vector"
    ),
    contains = "matrix"
)


###
# Nifti (nifti_array) Methods
###

setGeneric('is.nifti', 
    function(x) standardGeneric('is.nifti')
)

setMethod('is.nifti',
    signature(x='nifti_array'),
    function(x) return(TRUE)
)

setMethod('is.nifti',
    definition=function(x) return(FALSE)
)

setGeneric('as.nifti', 
    function(x, header=NULL) standardGeneric('as.nifti')
)

setMethod('as.nifti',
    signature(x='nifti', header='missing'),
    function(x) {
        ndim <- length(dim(x))
        if (ndim == 3)
            y <- x[,,]
        else if (ndim == 4)
            y <- x[,,,]
        else
            stop(sprintf('nifti dimension %i is not supported (only 3D and 4D)', ndim))
        return(as.nifti(y, getHeader(x)))
    }
)

setMethod('as.nifti',
    signature(x='nifti_array', header='missing'),
    function(x) return(x)
)

setMethod('as.nifti',
    signature(x='nifti_matrix', header='missing'),
    function(x) {        
        tmp <- unmask(x)
        dim(tmp) <- x@header$dim
        return(as.nifti(tmp, x@header))
    }
)

setMethod('as.nifti',
    signature(x='vector', header='nifti_matrix'),
    function(x, header) {    
        niftimat <- header
        header <- niftimat@header
        
        if (length(x) > length(niftimat@mask_used))
            stop('input vector too long...must be long enough to transform to 3D space')
        
        header$dim <- header$dim[1:3]
        header$pixdim <- header$pixdim[1:3]
        
        tmp <- vector(class(x), length(niftimat@mask_used))
        tmp[x@mask_used] <- x
        dim(tmp) <- header$dim
        
        return(as.nifti(tmp, header))
    }
)

setMethod('as.nifti',
    signature(x='array'),
    function(x, header) {
        arr <- new("nifti_array", header=header)
        arr@.Data <- x
        return(arr)
    }
)

read.nifti <- function(fname) {
    x <- read.nifti.old(fname)
    return(as.nifti(x))
}

read.nifti.header <- function(fname) {
    x <- read.nifti.old.header(fname)
    return(getHeader(x))
}

setGeneric('hdr', function(x) standardGeneric('hdr'))

setMethod('hdr',
    signature(x='nifti_array'),
    function(x) x@header
)

header <- function(x) hdr(x)

setGeneric('img', function(x) standardGeneric('img'))

setMethod('img',
    signature(x='nifti_array'),
    function(x) x@.Data
)

###
# Nifti2Mat (nifti_matrix) Methods
###

as.nifti2mat <- function(x) {
    if (!is.nifti(x))
        stop("input must be of type 'nifti_array'")
    if (length(dim(x)) != 4)
        stop("dimensions not equal to 4")
    
    # 1. Create 2D matrix
    tmpMat <- x@.Data
    dim(tmpMat) <- c(prod(dim(x)[1:3]), dim(x)[4])  # rows = voxels & cols = timepoints
    tmpMat <- as.matrix(tmpMat)
    mat <- new("nifti_matrix", header=x@header, mask_used=rep(TRUE, nrow(tmpMat)))
    mat@.Data <- tmpMat
    
    # 2. Set dimnames with 'x:y:z' for each column name
    indices <- lapply(dim(x)[1:3], seq) # list indices in each dimension
    coords <- expand.grid(indices)  # rows=voxels, cols=x,y,z
    names(dim(mat)) <- c("voxels", "timepoints")
    rn <- apply(coords, 1, paste, collapse=":")
    dimnames(mat) <- list(voxels=rn, timepoints=1:dim(x)[4])
    
    return(mat)
}

read.nifti2mat <- function(fname) {
    x <- read.nifti(fname)
    return(as.nifti2mat(x))
}

getmask <- function(m, thresh=0) {
    if (!is.null(m) && !is.vector(m)) {
        if (is.character(m))
            m <- read.nifti(m)
        
        if (is.array(m)) {
            if (length(dim(m)) != 3)
                stop("if m is a nifti file or array, dimensions must be 3")
            m <- as.vector(m>thresh)
        } else {
            stop("unrecognized type for input 'm': ", class(m))
        }
    }
    
    return(m)
}

mask <- function(x, m) {
    if (class(x) != "nifti_matrix")
        stop("Unrecognized class for input 'x': ", class(x))
    
    # 1. Get mask (convert to vector)
    m <- getmask(m)
    
    # 2. Handle if 'x' already masked
    if (!all(x@mask_used)) {
        ## can only accept logical
        if (!is.logical(m))
            stop("If nifti_matrix object is already masked, new mask must be a logical")
        
        ## mask size of filtered/masked data
        ## ensure length(m) == # of voxels in original nifti file
        if (length(m) == nrow(x)) {
            tmp <- x@mask_used
            tmp[tmp] <- m
            m <- tmp
        }
    }
    
    # 3. Set m to be a vector of logicals (if not already)
    if (is.logical(m)) {
        if (length(m) != length(x@mask_used))
            stop("mask vector is an incorrect length, must be equal to number of voxels in original nifti")
    }
    ## vector of column indices
    else if (is.integer(m)) {
        if (any(duplicated(m)))
            stop("cannot have duplicate column indices")
        if (max(abs(m)) > length(x@mask_used))
            stop("a column index is larger than size of matrix")
        tmp <- rep(FALSE, nrow(x))
        tmp[m] <- m
        m <- tmp
    }
    ## error
    else {
        stop("Unrecognized class for input 'm': ", class(m))
    }
    
    # 4. Mask
    x@mask_used <- m
    x@.Data <- x[m,]
    
    return(x)
}

unmask <- function(x, bg=0) {
    if (class(x) != 'nifti_matrix')
        stop("input must be of type 'nifti_matrix'")
    
    nifti_dim <- x@header$dim
    tmpMat <- matrix(bg, prod(nifti_dim[1:3]), nifti_dim[4])
    tmpMat[x@mask_used,] <- x@.Data
    x@.Data <- tmpMat
    x@mask_used <- rep(TRUE, length(x@mask_used))
    
    return(x)
}

setMethod('hdr',
    signature(x='nifti_matrix'),
    function(x) x@header
)

# TODO: not sure about this 'img' method
setMethod('img',
    signature(x='nifti_matrix'),
    function(x) x@.Data
)

setGeneric('coords', function(x) standardGeneric('coords'))

setMethod('coords',
    signature(x='nifti_matrix'),
    function(x) .coords(x@header, x@mask_used)
)

setMethod('coords',
    signature(x='nifti_array'),
    function(x) .coords(x@header, rep(TRUE, prod(dim(x)[1:3])))
)

.coords <- function(header, mask_used) {
    indices <- lapply(header$dim[1:3], seq) # list indices in each dimension
    coords <- expand.grid(indices)  # rows=voxels, cols=x,y,z
    coords <- cbind(coords, mask_used)    # cols=x,y,z,voxel_used
    dimnames(coords) <- list(voxels=which(mask_used==T), c("x", "y", "z", "voxel_used"))
    return(coords)
}

sort.coords <- function(coords) {
    if (ncol(coords) < 3)
        stop("input must have at least 3 columns (x, y, z coordinates)")

    x.vals <- unique(coords[,1])
    y.vals <- unique(coords[,2])
    
    # Sort X axis
    coords <- coords[order(coords[,1]),]
    for (i in x.vals) {
        this.x <- coords[,1] == i
        yzcoords <- coords[this.x,-1]
        
        # Sort Y axis
        yzcoords <- yzcoords[order(yzcoords[,1]),]
        for (j in y.vals) {
            this.y <- yzcoords[,1] == j
            zcoords <- yzcoords[this.y,-1,drop=F]
            
            # Sort Z axis
            yzcoords[this.y,-1] <- zcoords[order(zcoords[,1]),]
        }
        
        coords[this.x,-1] <- yzcoords
    }
    
    return(coords)
}

###
# Methods for batch processing
###

batch.nifti2mat <- function(inputs, m=NULL, verbose=F) {
    m <- getmask(m)
    
    n <- length(inputs)
    nifti_list <- vector("list", n)
    if (verbose) {
        pb <- create.progressbar(n)
        update.progressbar(pb, 0)
    }
    
    for (i in 1:n) {
        if (is.character(inputs[[i]]))
            x <- read.nifti2mat(inputs[[i]])
        else if (is.nifti(inputs[[i]]))
            x <- as.nifti2mat(inputs[[i]])
        else
            x <- inputs[[i]]
        
        if (class(x) != "nifti_matrix")
            stop("unrecognized type: ", class(x))
        
        if (!is.null(m))
            x <- mask(x, m)
        
        nifti_list[[i]] <- x
        
        if (verbose)
            update.progressbar(pb, i)
    }
    
    gc(verbose=F, reset=T)
    
    if (verbose)
        end.progressbar(pb)
    
    return(nifti_list)
}


# TODO:
#   - plot function
#   - print function?
#   - future: what if create simulated data in R

###
# Header related functions
###

ijk2xyx <- function(ijk, header, use.qform=T) {
    if (use.qform)
        xyz <- ijk %*% header$qto_xyz[1:3,1:3] + header$qto_xyz[1:3,4]
    else
        xyz <- ijk %*% header$sto_xyz[1:3,1:3] + header$sto_xyz[1:3,4]
    return(xyz)
}

matrix.ijk2xyz <- function(mat.ijk, header, use.qform=T) {
    mat.ijk <- as.matrix(mat.ijk)
    
    if (use.qform) {
        tmp <- matrix(rep(header$qto_xyz[1:3,4], nrow(mat.ijk)), nrow(mat.ijk), byrow=T)
        mat.xyz <- mat.ijk %*% header$qto_xyz[1:3,1:3] + tmp
    } else {
        tmp <- matrix(rep(header$sto_xyz[1:3,4], nrow(mat.ijk)), nrow(mat.ijk), byrow=T)
        mat.xyz <- mat.ijk %*% header$sto_xyz[1:3,1:3] + tmp
    }
    return(mat.xyz)
}

###
# Internal wrappers to Rniftilib functions and attributes
###

read.nifti.old <- function(fname) {
    nifti.image.read(abspath(fname))
}

read.nifti.old.header <- function(fname) {
    nifti.image.read(abspath(fname), read_data=0)
}

getHeader <- function(niftiobj) {
    if (class(niftiobj) != "nifti")
        stop("Input must be a nifti object")
    
    # Get names of header attributes/names
    header_attributes <- c("qto.xyz", "qto.ijk", "sto.xyz", "sto.ijk", "toffset", "descrip", "fname", "iname", "slice.duration", "qform.code", "sform.code", "quatern.b", "quatern.c", "quatern.d", "qoffset.x", "qoffset.y", "qoffset.z", "qfac", "pixdim", "nifti.type", "sizeof.hdr", "datatype", "scl.slope", "scl.inter", "qto_xyz", "qto_ijk", "sto_xyz", "sto_ijk", "dim")
    n <- length(header_attributes)
    
    # Store header attributes into a list
    hdr <- lapply(1:n, function(i) .Call("Rnifti_image_getattribute", niftiobj, header_attributes[i], PACKAGE="Rniftilib"))
    names(hdr) <- header_attributes
    
    return(hdr)
}
